import os
import tqdm 
import fire
import openai
import itertools


import numpy   as      np
from   tqdm    import tqdm
from   tools   import  *
from   typing  import  List, Union, Iterable, Dict

BASE_URL    = "http://36.111.148.95:8000/v1"
API_KEY     = "123456"
MODEL_NAME  = "chatglm3-6b"

DATA_PREFIX = "/home/yangxianpku/datasets/"
OUT_PATH    = f"./{MODEL_NAME}.jsonl"
client      = openai.OpenAI(api_key=API_KEY, base_url=BASE_URL)

def exec(datasets, format_tabs:bool=False, prompt_type:int=1, exec_path=OUT_PATH):
    samples = []
    origins = []
    pbar    = tqdm.tqdm(total=len(datasets))
    for task_id, task in datasets.items():
        prompt   = fix_indent(task['prompt'], indent=4) if format_tabs else task['prompt']
        question = coding_prompt(prompt, prompt_type)
        resp     = client.chat.completions.create(
                        model    = MODEL_NAME,
                        messages = [
                            {"role": "system", "content":  f"你是一个资深的程序员."},
                            {"role": "user",   "content":  question}
                        ],
                        temperature=0.6,
                    ).choices[0].message.content
        result  = dict(task_id=task_id, resp=resp)
        origins += [result]
        
        resp = filter_code(fix_indent(resp, indent=4), MODEL_NAME)
        pbar.update(1)
        result  = dict(task_id=task_id, resp=resp)
        samples += [result]
    write_jsonl(exec_path, samples)
    write_jsonl(f"./{MODEL_NAME}_native.jsonl", origins)




def eval(datasets, exec_path:str=OUT_PATH, k:str="1,10,100", num_worker:int=4, **kwargs):
    resps = stream_jsonl(exec_path)     # 模型返回的代码响应
    k = list(map(int, k.split(",")))

    futures = []
    resp_id = Counter()
    n_resp  = 0
    rets    = defaultdict(list)

    for resp in resps:
        pass
    

def main(format_tabs:bool=False, prompt_type:int=1, **kwargs):    
    datasets = read_problems(DATA_PREFIX)
    exec(datasets, format_tabs, prompt_type)
    # eval_ret = eval(datasets, exec_path=OUT_PATH, k="1,10,100", num_worker=4, **kwargs)

    # print(eval_ret["pass@1"])
    # print(eval_ret["pass@10"])
    # print(eval_ret["pass@100"])


if __name__ == "__main__":
    fire.Fire(main)